#!/usr/bin/env python3
import argparse
import random
import socket
import sys
from scapy.all import IP, TCP, UDP, Ether, Packet, bind_layers, get_if_hwaddr, get_if_list, sendp
from scapy.fields import BitField


def get_net_interface():
    """获取网络接口"""
    ifaces = [i for i in get_if_list() if "eth0" in i]
    return ifaces[0] if ifaces else None


class SourceRouteHeader(Packet):
    """源路由头部"""
    fields_desc = [
        BitField("bos", 0, 1),    # 栈起始标志
        BitField("port", 0, 15)   # 端口号
    ]


# 设置协议绑定
bind_layers(Ether, SourceRouteHeader, type=0x1234)
bind_layers(SourceRouteHeader, SourceRouteHeader, bos=0)
bind_layers(SourceRouteHeader, IP, bos=1)


def build_ipv4_packet(iface, addr, message):
    """创建IPv4转发数据包"""
    pkt = Ether(src=get_if_hwaddr(iface), dst='ff:ff:ff:ff:ff:ff')
    pkt = pkt / IP(dst=addr) / TCP(dport=1234, sport=random.randint(49152, 65535)) / message
    return pkt


def build_source_route_packet(iface, addr, message, ports):
    """创建源路由数据包"""
    pkt = Ether(src=get_if_hwaddr(iface), dst='ff:ff:ff:ff:ff:ff')
    
    # 添加源路由头部
    for port in ports:
        pkt = pkt / SourceRouteHeader(bos=0, port=port)
    
    # 设置最后一个头部的BOS标志
    if pkt.haslayer(SourceRouteHeader):
        pkt.getlayer(SourceRouteHeader, len(ports)).bos = 1
    
    pkt = pkt / IP(dst=addr) / UDP(dport=4321, sport=1234) / message
    return pkt


def main():
    """主函数"""
    print("IPv4转发模式: ./send.py <目标IP> <消息> --mode ip")
    print("源路由模式: ./send.py <目标IP> <消息>")
    
    if len(sys.argv) < 3:
        print('用法: ./send.py <目标IP> <消息> [--mode ip]')
        return 1

    parser = argparse.ArgumentParser()
    parser.add_argument('ip_addr', type=str, help="目标IP地址")
    parser.add_argument('message', type=str, help="要包含的消息")
    parser.add_argument('--mode', type=str, default=None, help="模式: 'ip' 表示IPv4转发")
    args = parser.parse_args()

    iface = get_net_interface()
    if not iface:
        print("找不到eth0接口")
        return 1

    addr = socket.gethostbyname(args.ip_addr)
    print(f"在接口 {iface} 上发送到 {addr}")

    if args.mode == "ip":
        # IPv4转发模式
        pkt = build_ipv4_packet(iface, addr, args.message)
        pkt.show2()
        sendp(pkt, iface=iface, verbose=False)
    elif args.mode is None:
        # 源路由模式(交互)
        while True:
            print()
            user_input = input('输入空格分隔的端口号(例如: "2 3 2 2 1") 或 "q" 退出: ')
            
            if user_input.lower() == "q":
                break
            
            try:
                ports = [int(p) for p in user_input.split()]
                pkt = build_source_route_packet(iface, addr, args.message, ports)
                pkt.show2()
                sendp(pkt, iface=iface, verbose=False)
            except ValueError:
                print("无效输入，请只输入数字")
    else:
        print("模式必须是'ip'或留空使用源路由")
        return 1

    return 0


if __name__ == '__main__':
    sys.exit(main())